from rainbowneko.data import VisionDataSource, ImageLabelSource
from rainbowneko.utils import Path_Like, is_image_file
from pathlib import Path
from typing import Dict, Any
import re
import numpy as np

class YoloSource(VisionDataSource):
    def __init__(self, img_root: Path_Like, label_root: Path_Like, repeat=1, **kwargs):
        super().__init__(img_root, repeat=repeat)
        self.label_root = Path(label_root)

        self.img_ids = self._load_img_ids(img_root)

    def _load_img_ids(self, img_root: Path_Like):
        return [x.name for x in Path(img_root).iterdir() if is_image_file(x)] * self.repeat

    def __len__(self):
        return len(self.img_ids)

    def load_label(self, img_id):
        path = self.label_root / (Path(img_id).stem+'.txt')
        bboxs = path.read_text(encoding='utf-8').splitlines()

        labels = []
        for line in bboxs:
            if len(line) > 0:
                items = re.split(r'\s+', line)
                cls, x, y, w, h = items
                x, y, w, h = float(x), float(y), float(w), float(h)
                labels.append([int(cls), x - w / 2, y - h / 2, x + w / 2, y + h / 2])
        return labels

    def __getitem__(self, index) -> Dict[str, Any]:
        img_id = self.img_ids[index]
        path = self.img_root / img_id
        return {
            'id': img_id,
            'image': path,
            'label': self.load_label(img_id)
        }